#include <stdlib.h>
#include <memory.h>
#include <stdio.h>
#include "htea.h"

#if _MSC_VER
#pragma warning(disable:4996)
#endif // #if _MSC_VER

#define TEA_DELTA                       0x9e3779b9
#define TEA_ROTATE                      32
#define TEA_MAX_BUFFSIZE                (4 * 1024 * 1024)
#define TEA_HEADER_CHECKER              0x47ae5f03
#define TEA_HDRSIZE                     sizeof(TEA_HEADER)

typedef unsigned int uint32_t;

typedef struct _TEA_HEADER
{
    unsigned int hdr;
    unsigned int cksum;
    int size;
    int size_high;
} TEA_HEADER, *PTEA_HEADER;

static void _tea_encrypt(uint32_t* v, uint32_t* k)
{
    uint32_t v0 = v[0], v1 = v[1], sum = 0, i;              /* set up */
    uint32_t k0 = k[0], k1 = k[1], k2 = k[2], k3 = k[3];    /* cache key */
    for (i = 0; i < TEA_ROTATE; i++)
    {
        sum += TEA_DELTA;
        v0 += ((v1 << 4) + k0) ^ (v1 + sum) ^ ((v1 >> 5) + k1);
        v1 += ((v0 << 4) + k2) ^ (v0 + sum) ^ ((v0 >> 5) + k3);
    }
    v[0] = v0;
    v[1] = v1;
}

static void _tea_decrypt(uint32_t* v, uint32_t* k)
{
    uint32_t v0 = v[0], v1 = v[1], sum = 0xC6EF3720, i;     /* set up */
    uint32_t k0 = k[0], k1 = k[1], k2 = k[2], k3 = k[3];    /* cache key */
    for (i = 0; i < TEA_ROTATE; i++)
    {
        v1 -= ((v0 << 4) + k2) ^ (v0 + sum) ^ ((v0 >> 5) + k3);
        v0 -= ((v1 << 4) + k0) ^ (v1 + sum) ^ ((v1 >> 5) + k1);
        sum -= TEA_DELTA;
    }
    v[0] = v0;
    v[1] = v1;
}

static int _tea_encode_file(FILE* fsrc, FILE* fdst, unsigned int key[4])
{
    int ret = 0;
    char read_buf[1024];
    char tmp_buf[1040];
    int file_size = 0;
    int readed = 0;
    int enced_size = 0;
    int state = 0;
    TEA_HEADER hdr = {0};

    do
    {
        fseek(fsrc, 0, SEEK_END);
        file_size = (int)ftell(fsrc);
        fseek(fsrc, 0, SEEK_SET);
        if (file_size > 0x7ffffff0 || (0 == file_size))
        {
            break;
        }

        hdr.hdr = rand() * rand();
        hdr.cksum = ~(hdr.hdr ^ TEA_HEADER_CHECKER);
        hdr.size = file_size;

        fwrite(&hdr, TEA_HDRSIZE, 1, fdst);

        while ((readed = (int)fread(read_buf, 1, 1024, fsrc)))
        {
            enced_size = 1040;
            if (0 != htea_enc_buff(read_buf, readed, key, tmp_buf, &enced_size))
            {
                state = 1;
                break;
            }

            if ((enced_size - TEA_HDRSIZE) != fwrite(tmp_buf + TEA_HDRSIZE, 1, enced_size - TEA_HDRSIZE, fdst))
            {
                state = 1;
                break;
            }
        }

        ret = state;
    } while (0);

    return ret;
}

static int _tea_decode_file(FILE* fsrc, FILE* fdst, unsigned int key[4])
{
    int ret = 0;
    int file_size = 0;
    int new_file_size = 0;
    int state = 0;
    TEA_HEADER hdr = {0};
    int i;
    int readed = 0;
    int need_write = 0;
    int total_writed = 0;
    char read_buf[1024];

    do
    {
        fseek(fsrc, 0, SEEK_END);
        file_size = (int)ftell(fsrc);
        fseek(fsrc, 0, SEEK_SET);
        if (file_size > 0x7ffffff0 || file_size < TEA_HDRSIZE)
        {
            break;
        }

        fread(&hdr, 1, TEA_HDRSIZE, fsrc);
        // verify file
        if (hdr.cksum != (~(hdr.hdr ^ TEA_HEADER_CHECKER)))
        {
            break;
        }
        new_file_size = hdr.size;

        while ((readed = (int)fread(read_buf, 1, 1024, fsrc)))
        {
            if (readed % TEA_UNITSIZE)
            {
                break;
            }

            for (i = 0; i < readed; i += TEA_UNITSIZE)
            {
                _tea_decrypt((uint32_t*)(read_buf + i), key);
            }

            need_write = (total_writed + readed) > new_file_size ? (new_file_size - total_writed) : readed;
            if (need_write != fwrite(read_buf, 1, need_write, fdst))
            {
                state = 1;
                break;
            }
            total_writed += need_write;
        }

        ret = state;
    } while (0);

    return ret;
}

static int _tea_file_handler(const char* src_file, const char* dst_file, unsigned int key[4], int encode)
{
    int ret = 1;
    FILE* fsrc = 0;
    FILE* fdst = 0;

    do
    {
        fsrc = fopen(src_file, "rb");
        if (!fsrc)
        {
            break;
        }

        fdst = fopen(dst_file, "wb");
        if (!fdst)
        {
            break;
        }

        if (encode)
        {
            ret = _tea_encode_file(fsrc, fdst, key);
        }
        else
        {
            ret = _tea_decode_file(fsrc, fdst, key);
        }
    } while (0);

    if (fsrc)
    {
        fclose(fsrc);
    }
    if (fdst)
    {
        fclose(fdst);
    }

    return ret;
}

int HPLATFORMCALL htea_enc_buff(const char* buf, int in_size, unsigned int key[4], char* out_buf, int* out_size)
{
    int need_len = 0;
    TEA_HEADER hdr = {0};
    int i;

    if (!buf || !in_size || !out_size || in_size > TEA_MAX_BUFFSIZE)
    {
        return 1;
    }

    need_len = TEA_HDRSIZE + in_size + ((in_size % TEA_UNITSIZE) ? (TEA_UNITSIZE - in_size % TEA_UNITSIZE) : 0);
    if (*out_size < need_len)
    {
        *out_size = need_len;
        return 2;
    }
    *out_size = need_len;

    hdr.hdr = rand() * rand();
    hdr.cksum = ~(hdr.hdr ^ TEA_HEADER_CHECKER);
    hdr.size = in_size;

    memcpy(out_buf, &hdr, TEA_HDRSIZE);
    memcpy(out_buf + TEA_HDRSIZE, buf, in_size);
    if (in_size % TEA_UNITSIZE)
    {
        memset(out_buf + TEA_HDRSIZE + in_size, 0, (TEA_UNITSIZE - in_size % TEA_UNITSIZE));
    }

    for (i = TEA_HDRSIZE; i < need_len; i += TEA_UNITSIZE)
    {
        _tea_encrypt((uint32_t*)(out_buf + i), key);
    }

    return 0;
}

int HPLATFORMCALL htea_dec_buff(const char* buf, int in_size, unsigned int key[4], char* out_buf, int* out_size)
{
    PTEA_HEADER pHdr = (PTEA_HEADER)buf;
    char* tmp = 0;
    unsigned int i = 0;

    if (!buf || !in_size || !out_size || in_size > TEA_MAX_BUFFSIZE || in_size < 24)
    {
        return 2;
    }

    if ((~(pHdr->hdr ^ TEA_HEADER_CHECKER)) != pHdr->cksum)
    {
        return 1;
    }

    if (*out_size < pHdr->size)
    {
        *out_size = pHdr->size;
        return 3;
    }
    *out_size = pHdr->size;

    tmp = (char*)malloc(in_size - TEA_HDRSIZE);
    if (!tmp)
    {
        return 4;
    }
    memcpy(tmp, buf + TEA_HDRSIZE, in_size - TEA_HDRSIZE);

    for (; i < in_size - TEA_HDRSIZE; i += TEA_UNITSIZE)
    {
        _tea_decrypt((uint32_t*)(tmp + i), key);
    }

    memcpy(out_buf, tmp, pHdr->size);
    free(tmp);

    return 0;
}

int HPLATFORMCALL htea_enc_file(const char* src_file, const char* dst_file, unsigned int key[4])
{
    return _tea_file_handler(src_file, dst_file, key, 1);
}

int HPLATFORMCALL htea_dec_file(const char* src_file, const char* dst_file, unsigned int key[4])
{
    return _tea_file_handler(src_file, dst_file, key, 0);
}

int HPLATFORMCALL htea_enc_buff_naked(const char* buf, int in_size, unsigned int key[4], char* out_buf, int* out_size)
{
    int group_count = in_size / TEA_UNITSIZE + (in_size % TEA_UNITSIZE ? 1 : 0);
    int idx = 0;

    if (!out_size || *out_size < group_count * TEA_UNITSIZE)
    {
        return 1;
    }

    *out_size = group_count * TEA_UNITSIZE;
    for (idx = 0; idx < group_count; ++idx)
    {
        int offset = idx * TEA_UNITSIZE;
        int buf_size = (in_size - offset >= 0) ? TEA_UNITSIZE : (offset - in_size);
        memcpy(out_buf + offset, buf + offset, buf_size);
        _tea_encrypt((uint32_t*)(out_buf + offset), key);
    }

    return 0;
}

int HPLATFORMCALL htea_dec_buff_naked(const char* buf, int in_size, unsigned int key[4], char* out_buf, int* out_size)
{
    int group_count = in_size / TEA_UNITSIZE + (in_size % TEA_UNITSIZE ? 1 : 0);
    int idx = 0;

    if (in_size % TEA_UNITSIZE || *out_size < in_size)
    {
        return 1;
    }
    *out_size = in_size;

    for (idx = 0; idx < group_count; ++idx)
    {
        int offset = idx * TEA_UNITSIZE;
        memcpy(out_buf + offset, buf + offset, TEA_UNITSIZE);
        _tea_decrypt((uint32_t*)(out_buf + offset), key);
    }

    return 0;
}
